package middleware

import (
	"gitee.com/igolang/imoney/internal/common/httpResp"
	"github.com/gin-gonic/gin"
	"net/http"
	"sync/atomic"
	"time"
)

type limit struct {
	ts    int64
	max   int64
	count int64
}

func NewLimit(max int) *limit {
	return &limit{max: int64(max), ts: time.Now().Unix()}
}

func (l *limit) OK() bool {
	tsOld := atomic.LoadInt64(&l.ts)
	ts := time.Now().Unix()
	if ts != tsOld {
		atomic.StoreInt64(&l.ts, ts)
		atomic.StoreInt64(&l.count, 1)
		return true
	}
	return atomic.AddInt64(&l.count, 1) < l.max
}

func Limited(max int) gin.HandlerFunc {
	if max < 1 {
		return func(c *gin.Context) {
			c.Next()
		}
	}

	l := NewLimit(max)
	return func(c *gin.Context) {
		if !l.OK() {
			r := httpResp.Response()
			r.Message = "abort"
			r.Code = http.StatusBadRequest
			httpResp.AbortJSON(c, r)
			return
		}
		c.Next()
	}
}

func Limited2(max int) gin.HandlerFunc {
	if max < 1 {
		return func(c *gin.Context) {
			c.Next()
		}
	}
	ch := make(chan struct{}, max)
	return func(c *gin.Context) {
		if len(ch) >= max {
			r := httpResp.Response()
			r.Message = "abort"
			r.Code = http.StatusBadRequest
			httpResp.AbortJSON(c, r)
		}

		ch <- struct{}{}
		c.Next()
		<-ch
	}
}
